import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
# torch_npu.npu.set_compile_mode(jit_compile=True)

from lerobot.policies.pi0.modeling_pi0 import PI0Policy
from lerobot.configs.types import FeatureType, PolicyFeature

import time


def main():

    dtype = torch.float16
    print("===Loading model===")
    policy = PI0Policy.from_pretrained("/path/to/weights").to(device="npu")
    print("===Done===")

    policy.config.output_features = {
        "action": PolicyFeature(
            type=FeatureType.ACTION,
            shape=(32,)  
        )
    }

    policy.config.input_features = {
        "observation.images.top": PolicyFeature(
            type=FeatureType.VISUAL, 
            shape=(1, 3, 224, 224)
        ),
        "observation.images.left": PolicyFeature(
            type=FeatureType.VISUAL, 
            shape=(1, 3, 224, 224)
        ),
        "observation.images.right": PolicyFeature(
            type=FeatureType.VISUAL, 
            shape=(1, 3, 224, 224)
        ),
        "observation.state": PolicyFeature(
            type=FeatureType.STATE,
            shape=(1, 32) 
        )
    }

    with torch.no_grad():
        policy.eval()
        times = []
        for i in range(10):
            policy.reset()
            state = torch.randn(1, 32, dtype=dtype).npu()
            img_top = torch.rand(1, 3, 244, 244, dtype=dtype).npu()
            img_left = torch.rand(1, 3, 244, 244, dtype=dtype).npu()
            img_right = torch.rand(1, 3, 244, 244, dtype=dtype).npu()
            observation = {
                            "observation.state": state,
                            "observation.images.top": img_top,
                            "observation.images.left": img_left,
                            "observation.images.right": img_right,
                            "task": ["Do something."]
                            }
            start = time.perf_counter()
            action = policy.select_action(observation)
            end = time.perf_counter() - start
            print(f"===Inference {i} cost time: {end * 1000} ms.===")
            print(action.cpu())
            
            times.append(end * 1000)

            # profiler.step()
    
    times = times[1:]
    print(f"===Average cost time: {round(sum(times) / len(times), 2)} ms.")


if __name__ == '__main__':

    main()